from typing import Dict, Any, List
from .base_agent import BaseAgent
from utils import logger

class SecurityCheckerAgent(BaseAgent):
    """安全检查代理 - 负责代码安全审查和漏洞检测"""
    
    def __init__(self):
        super().__init__(
            name="安全专家",
            role="信息安全专家和安全审计师"
        )
    
    def _get_system_prompt(self) -> str:
        return """你是一个资深的信息安全专家，具有以下专业技能：

🔒 **安全领域**：
- Web应用安全（OWASP Top 10）
- 代码安全审计和漏洞检测
- 输入验证和数据sanitization
- 认证和授权机制设计
- 加密和数据保护

🛡️ **漏洞类型**：
- SQL注入和NoSQL注入
- 跨站脚本攻击（XSS）
- 跨站请求伪造（CSRF）
- 命令注入和代码注入
- 文件上传和路径遍历漏洞

🔍 **安全检查**：
- 输入验证完整性
- 输出编码和转义
- 权限控制和访问管理
- 敏感信息泄露检查
- 错误处理安全性

⚡ **安全最佳实践**：
- 最小权限原则
- 深度防御策略
- 安全编码规范
- 加密算法选择
- 安全配置管理

📊 **风险评估**：
- 漏洞严重性评级
- 利用难度评估
- 业务影响分析
- 修复优先级排序
- 安全测试建议

请提供专业、全面的安全审查和改进建议。用中文回复。"""
    
    def process(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """进行安全审查和漏洞检测"""
        generated_code = input_data.get('generated_code', '')
        code_files = input_data.get('code_files', {})
        architecture_design = input_data.get('architecture_design', '')
        
        if not generated_code and not code_files:
            return {
                'success': False,
                'error': '未提供待检查的代码'
            }
        
        # 如果没有解析的代码文件，尝试从生成的代码中解析
        if not code_files and generated_code:
            code_files = {'main.py': generated_code}
        
        # 进行安全分析
        security_analysis = self._analyze_security_issues(code_files)
        
        # 构造安全审查提示
        prompt = f"""
基于以下信息进行全面的安全审查：

**代码文件**：
{self._format_code_files(code_files)}

**架构设计**：
{architecture_design[:500] + '...' if len(architecture_design) > 500 else architecture_design}

**初步安全分析**：
{security_analysis}

**请进行全面的安全审查，包括**：

## 1. 漏洞检测
按OWASP Top 10分类检查：
- A01: 访问控制失效
- A02: 加密失效
- A03: 注入攻击
- A04: 不安全设计
- A05: 安全配置错误
- A06: 易受攻击的组件
- A07: 身份验证和会话管理失效
- A08: 软件和数据完整性失效
- A09: 安全日志和监控失效
- A10: 服务端请求伪造（SSRF）

## 2. 输入验证审查
- 用户输入验证完整性
- SQL注入防护检查
- XSS攻击防护检查
- 命令注入防护检查
- 文件上传安全检查

## 3. 认证和授权审查
- 认证机制安全性
- 会话管理安全性
- 权限控制完整性
- 密码策略检查
- 令牌管理安全性

## 4. 数据保护审查
- 敏感数据识别
- 加密算法选择
- 数据传输安全
- 数据存储安全
- 隐私保护措施

## 5. 安全配置审查
- 错误处理安全性
- 日志记录安全性
- 配置文件安全性
- 依赖库安全性
- 部署配置安全性

## 6. 风险评估和修复建议
按严重程度分级：
- 🔴 严重（Critical）：立即修复
- 🟠 高危（High）：优先修复
- 🟡 中危（Medium）：计划修复
- 🟢 低危（Low）：建议修复

为每个问题提供：
- 漏洞描述
- 利用场景
- 业务影响
- 修复方案
- 代码示例

请提供详细的安全审查报告。
"""
        
        try:
            security_report = self._call_llm(prompt)
            
            # 生成安全改进代码
            secure_code_suggestions = self._generate_secure_code(code_files, security_analysis)
            
            # 生成安全摘要
            security_summary = self._generate_security_summary(security_analysis, security_report)
            
            return {
                'success': True,
                'security_report': security_report,
                'security_analysis': security_analysis,
                'secure_code_suggestions': secure_code_suggestions,
                'security_summary': security_summary,
                'summary': f'安全审查完成，发现{security_summary["total_issues"]}个安全问题'
            }
            
        except Exception as e:
            logger.error(f"安全审查失败: {e}")
            return {
                'success': False,
                'error': str(e),
                'summary': '安全审查失败'
            }
    
    def _analyze_security_issues(self, code_files: Dict[str, str]) -> Dict[str, Any]:
        """分析代码中的安全问题"""
        analysis = {
            'injection_risks': [],
            'xss_risks': [],
            'authentication_issues': [],
            'data_exposure_risks': [],
            'file_operation_risks': [],
            'crypto_issues': [],
            'configuration_issues': []
        }
        
        for file_path, code_content in code_files.items():
            # 检测SQL注入风险
            injection_risks = self._detect_injection_risks(code_content, file_path)
            analysis['injection_risks'].extend(injection_risks)
            
            # 检测XSS风险
            xss_risks = self._detect_xss_risks(code_content, file_path)
            analysis['xss_risks'].extend(xss_risks)
            
            # 检测认证问题
            auth_issues = self._detect_authentication_issues(code_content, file_path)
            analysis['authentication_issues'].extend(auth_issues)
            
            # 检测数据暴露风险
            data_risks = self._detect_data_exposure(code_content, file_path)
            analysis['data_exposure_risks'].extend(data_risks)
            
            # 检测文件操作风险
            file_risks = self._detect_file_operation_risks(code_content, file_path)
            analysis['file_operation_risks'].extend(file_risks)
            
            # 检测加密问题
            crypto_issues = self._detect_crypto_issues(code_content, file_path)
            analysis['crypto_issues'].extend(crypto_issues)
        
        return analysis
    
    def _detect_injection_risks(self, code_content: str, file_path: str) -> List[Dict[str, Any]]:
        """检测注入攻击风险"""
        risks = []
        lines = code_content.split('\n')
        
        # SQL注入风险模式
        sql_patterns = [
            'execute(',
            'cursor.execute(',
            'query =',
            'sql =',
            'SELECT',
            'INSERT',
            'UPDATE',
            'DELETE'
        ]
        
        for i, line in enumerate(lines):
            line_lower = line.lower()
            
            # 检测可能的SQL注入
            if any(pattern.lower() in line_lower for pattern in sql_patterns):
                if ('format(' in line or '%' in line or 'f"' in line) and 'user' in line_lower:
                    risks.append({
                        'type': 'sql_injection',
                        'file': file_path,
                        'line': i + 1,
                        'severity': 'high',
                        'description': '检测到可能的SQL注入风险：直接拼接用户输入'
                    })
            
            # 检测命令注入
            if ('os.system(' in line or 'subprocess.call(' in line or 'exec(' in line):
                if ('input(' in line or 'request.' in line or 'form.' in line):
                    risks.append({
                        'type': 'command_injection',
                        'file': file_path,
                        'line': i + 1,
                        'severity': 'critical',
                        'description': '检测到命令注入风险：执行用户可控的命令'
                    })
        
        return risks
    
    def _detect_xss_risks(self, code_content: str, file_path: str) -> List[Dict[str, Any]]:
        """检测XSS攻击风险"""
        risks = []
        lines = code_content.split('\n')
        
        for i, line in enumerate(lines):
            # 检测未转义的输出
            if 'render_template(' in line or 'return ' in line:
                if ('user' in line.lower() or 'input' in line.lower()) and 'escape(' not in line:
                    risks.append({
                        'type': 'xss',
                        'file': file_path,
                        'line': i + 1,
                        'severity': 'medium',
                        'description': '检测到XSS风险：未转义的用户输入输出'
                    })
            
            # 检测innerHTML等危险操作
            if ('innerHTML' in line or 'document.write(' in line):
                risks.append({
                    'type': 'dom_xss',
                    'file': file_path,
                    'line': i + 1,
                    'severity': 'high',
                    'description': '检测到DOM XSS风险：直接操作DOM'
                })
        
        return risks
    
    def _detect_authentication_issues(self, code_content: str, file_path: str) -> List[Dict[str, Any]]:
        """检测认证相关问题"""
        issues = []
        lines = code_content.split('\n')
        
        for i, line in enumerate(lines):
            # 检测硬编码密码
            if ('password' in line.lower() and ('=' in line or ':' in line)):
                if ('"' in line or "'" in line) and not line.strip().startswith('#'):
                    issues.append({
                        'type': 'hardcoded_password',
                        'file': file_path,
                        'line': i + 1,
                        'severity': 'high',
                        'description': '检测到硬编码密码风险'
                    })
            
            # 检测弱密码验证
            if 'len(password)' in line and any(num in line for num in ['<', '6', '4', '8']):
                issues.append({
                    'type': 'weak_password_policy',
                    'file': file_path,
                    'line': i + 1,
                    'severity': 'medium',
                    'description': '检测到弱密码策略'
                })
            
            # 检测不安全的会话管理
            if 'session[' in line and 'secure=False' in line:
                issues.append({
                    'type': 'insecure_session',
                    'file': file_path,
                    'line': i + 1,
                    'severity': 'medium',
                    'description': '检测到不安全的会话配置'
                })
        
        return issues
    
    def _detect_data_exposure(self, code_content: str, file_path: str) -> List[Dict[str, Any]]:
        """检测数据暴露风险"""
        risks = []
        lines = code_content.split('\n')
        
        # 敏感信息关键词
        sensitive_keywords = ['password', 'secret', 'key', 'token', 'api_key', 'private']
        
        for i, line in enumerate(lines):
            line_lower = line.lower()
            
            # 检测敏感信息日志
            if 'print(' in line or 'log' in line_lower:
                if any(keyword in line_lower for keyword in sensitive_keywords):
                    risks.append({
                        'type': 'sensitive_data_logging',
                        'file': file_path,
                        'line': i + 1,
                        'severity': 'medium',
                        'description': '检测到敏感信息可能被记录到日志'
                    })
            
            # 检测调试信息泄露
            if 'debug=True' in line:
                risks.append({
                    'type': 'debug_mode',
                    'file': file_path,
                    'line': i + 1,
                    'severity': 'low',
                    'description': '检测到调试模式开启，可能泄露敏感信息'
                })
        
        return risks
    
    def _detect_file_operation_risks(self, code_content: str, file_path: str) -> List[Dict[str, Any]]:
        """检测文件操作风险"""
        risks = []
        lines = code_content.split('\n')
        
        for i, line in enumerate(lines):
            # 检测路径遍历风险
            if 'open(' in line and ('input(' in line or 'request.' in line):
                if '../' not in line or 'os.path.join(' not in line:
                    risks.append({
                        'type': 'path_traversal',
                        'file': file_path,
                        'line': i + 1,
                        'severity': 'high',
                        'description': '检测到路径遍历风险：未验证的文件路径'
                    })
            
            # 检测文件上传风险
            if 'save(' in line and 'file' in line.lower():
                if 'allowed_extensions' not in code_content:
                    risks.append({
                        'type': 'unsafe_file_upload',
                        'file': file_path,
                        'line': i + 1,
                        'severity': 'high',
                        'description': '检测到不安全的文件上传：缺少文件类型验证'
                    })
        
        return risks
    
    def _detect_crypto_issues(self, code_content: str, file_path: str) -> List[Dict[str, Any]]:
        """检测加密相关问题"""
        issues = []
        lines = code_content.split('\n')
        
        # 弱加密算法
        weak_crypto = ['md5', 'sha1', 'des', 'rc4']
        
        for i, line in enumerate(lines):
            line_lower = line.lower()
            
            # 检测弱加密算法
            if any(weak in line_lower for weak in weak_crypto):
                issues.append({
                    'type': 'weak_crypto',
                    'file': file_path,
                    'line': i + 1,
                    'severity': 'medium',
                    'description': f'检测到弱加密算法使用'
                })
            
            # 检测不安全的随机数生成
            if 'random.random()' in line or 'random.randint(' in line:
                if 'crypto' not in code_content.lower():
                    issues.append({
                        'type': 'weak_random',
                        'file': file_path,
                        'line': i + 1,
                        'severity': 'low',
                        'description': '检测到不安全的随机数生成，建议使用加密安全的随机数'
                    })
        
        return issues
    
    def _format_code_files(self, code_files: Dict[str, str]) -> str:
        """格式化代码文件为字符串"""
        formatted = ""
        for file_path, content in code_files.items():
            formatted += f"\n--- {file_path} ---\n"
            # 只显示前800个字符以避免过长
            content_preview = content[:800]
            if len(content) > 800:
                content_preview += f"\n... [文件内容已截断，总长度: {len(content)} 字符]"
            formatted += content_preview + "\n"
        return formatted
    
    def _generate_secure_code(self, code_files: Dict[str, str], security_analysis: Dict[str, Any]) -> str:
        """生成安全改进代码建议"""
        # 收集所有安全问题
        all_issues = []
        for issue_type, issues in security_analysis.items():
            all_issues.extend(issues)
        
        if not all_issues:
            return "未发现明显的安全问题，代码安全性良好。"
        
        prompt = f"""
基于以下安全分析结果，提供安全改进的代码示例：

**发现的安全问题**：
{str(all_issues)[:1000]}

**原始代码文件**：
{self._format_code_files(code_files)[:1000]}

请提供：
1. 针对每个安全问题的修复代码示例
2. 安全最佳实践的实现
3. 输入验证和sanitization的代码
4. 安全配置的建议

用代码注释说明安全改进的原因。
"""
        
        return self._call_llm(prompt)
    
    def _generate_security_summary(self, security_analysis: Dict[str, Any], security_report: str) -> Dict[str, Any]:
        """生成安全摘要"""
        total_issues = sum(len(issues) for issues in security_analysis.values())
        
        # 按严重程度统计
        critical_count = security_report.count('🔴') + security_report.count('critical')
        high_count = security_report.count('🟠') + security_report.count('high')
        medium_count = security_report.count('🟡') + security_report.count('medium')
        low_count = security_report.count('🟢') + security_report.count('low')
        
        # 安全风险评级
        if critical_count > 0:
            risk_level = 'critical'
        elif high_count > 0:
            risk_level = 'high'
        elif medium_count > 0:
            risk_level = 'medium'
        else:
            risk_level = 'low'
        
        return {
            'total_issues': total_issues,
            'critical_issues': critical_count,
            'high_issues': high_count,
            'medium_issues': medium_count,
            'low_issues': low_count,
            'risk_level': risk_level,
            'injection_risks': len(security_analysis.get('injection_risks', [])),
            'xss_risks': len(security_analysis.get('xss_risks', [])),
            'authentication_issues': len(security_analysis.get('authentication_issues', [])),
            'data_exposure_risks': len(security_analysis.get('data_exposure_risks', [])),
            'file_operation_risks': len(security_analysis.get('file_operation_risks', [])),
            'crypto_issues': len(security_analysis.get('crypto_issues', []))
        }
    
    def generate_security_checklist(self, project_type: str = 'web') -> str:
        """生成安全检查清单"""
        prompt = f"""
为{project_type}项目生成详细的安全检查清单：

请包含：
1. 开发阶段安全检查项
2. 测试阶段安全验证项
3. 部署阶段安全配置项
4. 运维阶段安全监控项
5. 定期安全审计项

每个检查项包含检查方法和验收标准。
"""
        
        return self._call_llm(prompt) 